#!/usr/bin/python3
# coding=utf8
from __future__ import print_function, division, absolute_import

import copy
import threading
import time
import signal
import sys

import open3d as o3d
import rospy
import ros_numpy
from geometry_msgs.msg import PoseWithCovarianceStamped, Pose, Point, Quaternion, TransformStamped
from nav_msgs.msg import Odometry
from sensor_msgs.msg import PointCloud2
import numpy as np
import tf
import tf.transformations
import tf2_ros

from std_msgs.msg import String

# ================== Global Variables ==================
global_map = None
last_raw_map_points_num = 0
initialized = False
T_map_to_odom = np.eye(4)  # 默认初始位姿为 (0,0,0)
cur_odom = None
cur_scan = None

global_map_initialized = False

# Threading control
loc_thread = None
shutdown_flag = threading.Event()

# TF broadcaster
tf_broadcaster = None

# Publishers
pub_pc_in_map = None
pub_submap = None
pub_map_to_odom = None

# Parameters (set in main)
MAP_VOXEL_SIZE = 0.3
SCAN_VOXEL_SIZE = 0.1
FREQ_LOCALIZATION = 1
LOCALIZATION_TH = 0.80
FOV = 2 * 3.1415926
FOV_FAR = 30


# ================== Helper Functions ==================

def pose_to_mat(pose_msg):
    """Convert geometry_msgs/Pose to 4x4 SE(3) matrix."""
    pos = pose_msg.pose.pose.position
    quat = pose_msg.pose.pose.orientation
    trans = tf.transformations.translation_matrix([pos.x, pos.y, pos.z])
    rot = tf.transformations.quaternion_matrix([quat.x, quat.y, quat.z, quat.w])
    return np.dot(trans, rot)


def msg_to_array(pc_msg):
    pc_array = ros_numpy.numpify(pc_msg)
    pc = np.zeros([len(pc_array), 3])
    pc[:, 0] = pc_array['x']
    pc[:, 1] = pc_array['y']
    pc[:, 2] = pc_array['z']
    return pc


def voxel_down_sample(pcd, voxel_size):
    try:
        return pcd.voxel_down_sample(voxel_size)
    except AttributeError:
        return o3d.geometry.voxel_down_sample(pcd, voxel_size)


def inverse_se3(trans):
    trans_inv = np.eye(4)
    trans_inv[:3, :3] = trans[:3, :3].T
    trans_inv[:3, 3] = -trans[:3, :3].T @ trans[:3, 3]
    return trans_inv


def publish_point_cloud(publisher, header, pc):
    if len(pc) == 0:
        return
    data = np.zeros(len(pc), dtype=[
        ('x', np.float32),
        ('y', np.float32),
        ('z', np.float32),
        ('intensity', np.float32),
    ])
    data['x'] = pc[:, 0]
    data['y'] = pc[:, 1]
    data['z'] = pc[:, 2]
    data['intensity'] = 0.0
    msg = ros_numpy.msgify(PointCloud2, data)
    msg.header = header
    publisher.publish(msg)


def crop_global_map_in_FOV(global_map, pose_estimation, cur_odom):
    T_odom_to_base_link = pose_to_mat(cur_odom)
    T_map_to_base_link = pose_estimation @ T_odom_to_base_link
    T_base_link_to_map = inverse_se3(T_map_to_base_link)

    global_map_points = np.asarray(global_map.points)
    if global_map_points.size == 0:
        empty_pcd = o3d.geometry.PointCloud()
        return empty_pcd

    pts_hom = np.hstack([global_map_points, np.ones((len(global_map_points), 1))])
    pts_in_base = (T_base_link_to_map @ pts_hom.T).T

    x = pts_in_base[:, 0]
    y = pts_in_base[:, 1]
    dist_sq = x**2 + y**2
    angles = np.abs(np.arctan2(y, x))

    if FOV >= 2 * np.pi - 0.1:
        mask = dist_sq < FOV_FAR**2
    else:
        mask = (x > 0) & (dist_sq < FOV_FAR**2) & (angles < FOV / 2.0)

    selected = global_map_points[mask]
    cropped = o3d.geometry.PointCloud()
    cropped.points = o3d.utility.Vector3dVector(selected)

    header = copy.deepcopy(cur_odom.header)
    header.frame_id = 'map'
    publish_point_cloud(pub_submap, header, selected[::10])
    return cropped


def registration_at_scale(pc_scan, pc_map, initial, scale):
    scan_ds = voxel_down_sample(pc_scan, SCAN_VOXEL_SIZE * scale)
    map_ds = voxel_down_sample(pc_map, MAP_VOXEL_SIZE * scale)
    result = o3d.pipelines.registration.registration_icp(
        scan_ds, map_ds,
        1.0 * scale, initial,
        o3d.pipelines.registration.TransformationEstimationPointToPoint(),
        o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=20)
    )
    return result.transformation, result.fitness


def broadcast_tf():
    """Broadcast map -> odom transform."""
    global T_map_to_odom, tf_broadcaster
    t = TransformStamped()
    t.header.stamp = rospy.Time.now()  # 或使用 cur_odom 时间戳更精确
    t.header.frame_id = "map"
    t.child_frame_id = "odom"
    t.transform.translation.x = T_map_to_odom[0, 3]
    t.transform.translation.y = T_map_to_odom[1, 3]
    t.transform.translation.z = T_map_to_odom[2, 3]
    quat = tf.transformations.quaternion_from_matrix(T_map_to_odom)
    t.transform.rotation.x = quat[0]
    t.transform.rotation.y = quat[1]
    t.transform.rotation.z = quat[2]
    t.transform.rotation.w = quat[3]
    tf_broadcaster.sendTransform(t)


def global_localization(pose_estimation):
    global global_map, cur_scan, cur_odom, T_map_to_odom
    if cur_scan is None or cur_odom is None or global_map is None:
        return False

    rospy.loginfo('Global localization by scan-to-map matching...')

    scan_tobe_mapped = copy.deepcopy(cur_scan)
    global_map_copy = copy.deepcopy(global_map)

    tic = time.time()
    global_map_in_FOV = crop_global_map_in_FOV(global_map_copy, pose_estimation, cur_odom)
    if len(global_map_in_FOV.points) == 0:
        rospy.logwarn("No points in FOV!")
        return False

    transformation, _ = registration_at_scale(scan_tobe_mapped, global_map_in_FOV, initial=pose_estimation, scale=5)
    transformation, fitness = registration_at_scale(scan_tobe_mapped, global_map_in_FOV, initial=transformation, scale=1)

    toc = time.time()
    rospy.loginfo(f'ICP Time: {toc - tic:.3f}s, Fitness: {fitness:.3f}')

    if fitness > LOCALIZATION_TH:
        T_map_to_odom = transformation

        # Publish Odometry (for debugging)
        map_to_odom_msg = Odometry()
        xyz = tf.transformations.translation_from_matrix(T_map_to_odom)
        quat = tf.transformations.quaternion_from_matrix(T_map_to_odom)
        map_to_odom_msg.pose.pose = Pose(Point(*xyz), Quaternion(*quat))
        map_to_odom_msg.header.stamp = cur_odom.header.stamp
        map_to_odom_msg.header.frame_id = 'map'
        map_to_odom_msg.child_frame_id = 'odom'
        pub_map_to_odom.publish(map_to_odom_msg)

        # Broadcast TF for navigation stack
        broadcast_tf()
        return True
    else:
        rospy.logwarn(f"Localization failed! Fitness: {fitness:.3f}")
        return False


# ================== Callbacks ==================

def cb_save_cur_odom(odom_msg):
    global cur_odom
    cur_odom = odom_msg


def cb_save_cur_scan(pc_msg):
    global cur_scan
    # Do NOT override frame_id or stamp — keep original from Fast-LIO
    pub_pc_in_map.publish(pc_msg)

    try:
        # Fix field order if needed (Fast-LIO known issue)
        if len(pc_msg.fields) >= 8:
            pc_msg.fields = [
                pc_msg.fields[0], pc_msg.fields[1], pc_msg.fields[2],
                pc_msg.fields[4], pc_msg.fields[5], pc_msg.fields[6],
                pc_msg.fields[3], pc_msg.fields[7]
            ]
        pc = msg_to_array(pc_msg)
        cur_scan = o3d.geometry.PointCloud()
        cur_scan.points = o3d.utility.Vector3dVector(pc[:, :3])
    except Exception as e:
        rospy.logerr(f"Failed to process scan: {e}")


def initialize_global_map(pc_msg):
    global global_map, global_map_initialized, last_raw_map_points_num
    try:
        new_points = msg_to_array(pc_msg)[:, :3]
        new_points_num = new_points.shape[0]
        if 'last_raw_map_points_num' in globals() and last_raw_map_points_num == new_points_num:
            rospy.loginfo('Global map unchanged, skip update.')
            return
        last_raw_map_points_num = new_points_num

        global_map = o3d.geometry.PointCloud()
        global_map.points = o3d.utility.Vector3dVector(new_points)
        global_map = voxel_down_sample(global_map, MAP_VOXEL_SIZE)
        rospy.loginfo(f'Global map received with {len(global_map.points)} points.')
        global_map_initialized = True
    except Exception as e:
        rospy.logerr(f"Error initializing global map: {e}")


def cb_initial_pose(pose_msg):
    """Dynamically reset T_map_to_odom when /initialpose is published."""
    global T_map_to_odom
    try:
        T_map_to_odom = pose_to_mat(pose_msg)
        rospy.loginfo("Initial pose updated dynamically.")
        # Optional: trigger immediate localization
        # global_localization(T_map_to_odom)
    except Exception as e:
        rospy.logerr(f"Error updating initial pose: {e}")


# ================== Localization Thread ==================

def thread_localization():
    rate = rospy.Rate(FREQ_LOCALIZATION)
    while not rospy.is_shutdown() and not shutdown_flag.is_set():
        if global_map_initialized:
            global_localization(T_map_to_odom)
        rate.sleep()


# ================== Cleanup & Signal Handling ==================

def cleanup():
    rospy.loginfo("Shutting down localization node...")
    shutdown_flag.set()
    if loc_thread and loc_thread.is_alive():
        loc_thread.join(timeout=1.0)
    rospy.loginfo("Localization thread cleaned up.")


def signal_handler(sig, frame):
    cleanup()
    sys.exit(0)


# ================== Main ==================

if __name__ == '__main__':
    # Handle Ctrl+C gracefully
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    rospy.init_node('fast_lio_localization')
    rospy.loginfo('Localization Node Inited...')

    # Initialize TF broadcaster
    tf_broadcaster = tf2_ros.TransformBroadcaster()

    # Publishers
    pub_pc_in_map = rospy.Publisher('/cur_scan_in_map', PointCloud2, queue_size=1)
    pub_submap = rospy.Publisher('/submap', PointCloud2, queue_size=1)
    pub_map_to_odom = rospy.Publisher('/map_to_odom', Odometry, queue_size=1)

    # Subscribers
    rospy.Subscriber('/cloud_registered', PointCloud2, cb_save_cur_scan, queue_size=1)
    rospy.Subscriber('/Odometry', Odometry, cb_save_cur_odom, queue_size=1)
    rospy.Subscriber('/map3d', PointCloud2, initialize_global_map, queue_size=1)
    rospy.Subscriber('/initialpose', PoseWithCovarianceStamped, cb_initial_pose, queue_size=1)

    # Wait only for global map (no need to wait for initial pose)
    rospy.loginfo("Waiting for global map...")
    while not rospy.is_shutdown() and not global_map_initialized:
        rospy.sleep(0.2)

    if rospy.is_shutdown():
        sys.exit(0)

    rospy.loginfo("Global map ready. Starting localization with default pose (0,0,0).")

    # Start localization thread
    loc_thread = threading.Thread(target=thread_localization)
    loc_thread.daemon = True
    loc_thread.start()

    try:
        rospy.spin()
    finally:
        cleanup()